import json
import re
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    ServiceContext,
    load_index_from_storage,
    Document,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import (
    RetrieverQueryEngine,
    SubQuestionQueryEngine,
)
from llama_index.core.response_synthesizers import (
    ResponseMode,
    get_response_synthesizer,
)
from llama_index.core.tools import QueryEngineTool, ToolMetadata


from llama_index.vector_stores.elasticsearch import ElasticsearchStore

es = ElasticsearchStore(
    index_name="question_index",
    es_url="http://10.6.26.37:9200",
    es_user="elastic",
    es_password="fortune@123",
)

try:
    with open('data1-800.json', 'r', encoding='utf-8') as file:
        data = json.load(file)

    documents = []
    # 遍历每本书的信息
    for book in data:
        # 获取相关提问字段的值
        related_questions = book.get('相关提问', '')
        author = book.get('作者','无名')
        book_name = book.get('书名','')
        book_id = book.get('书籍id','')
        # 使用正则表达式匹配问题及答案'
        if related_questions:
            qa_pairs = re.split(r'\n', related_questions)
            if qa_pairs:
                for qa in qa_pairs:
                    qa = re.sub(r'^\d+\.\s*问：', '', qa)
                    document =[Document(text=qa,metadata={"file_name": book_name, "author": author})] 
                    documents.extend(document)

    storage_context = StorageContext.from_defaults(vector_store=es)

    index = VectorStoreIndex.from_documents(documents,storage_context=storage_context)

except Exception as e:
    print(f"An error occurred: {e}")
  
finally:
    if es:
        es.close()


  
    